package gp;

import java.io.PrintStream;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.SortedMap;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.Vector;

import tree.Tree;
import utils.Utils;

public class GeneticProgram 
{

	
	private int populationSize;
	private ArrayList<Tree> currentGeneration, nextGeneration;
	private int maxDepth;
	private int elitism, numRandomTreePerGen, numTreesForTournamentSelection, numNodesPenalty;
	private double mutationProbability, maxFitnessForElitism;
	
	private double bestFitness, averageFitness;
	
	public GeneticProgram(int populationSize, int maxDepth, double mutationProbability, int elitism,
			double maxFitnessForElitism, int numRandomTreePerGen, int numTreesForTournamentSelection, 
			int numNodesPenalty)
	{
		this.populationSize = populationSize;
		this.maxDepth = maxDepth;
		this.mutationProbability = mutationProbability;
		this.elitism = elitism;
		this.maxFitnessForElitism = maxFitnessForElitism;
		this.numRandomTreePerGen = numRandomTreePerGen;
		this.numTreesForTournamentSelection = numTreesForTournamentSelection;
		this.numNodesPenalty = numNodesPenalty;
		
		currentGeneration = new ArrayList<Tree>(populationSize);
		
		for (int i = 0; i < populationSize; i++)
		{
			currentGeneration.add(new Tree(maxDepth));
		}
	}
	
	/**
	 * train the next generation according to the training examples.
	 * Will go over all the training examples, so it's recommended to pass only
	 * a subset of them
	 * @param trainingExamples - an n by 2 array of training examples. 
	 * @param isLastStep TODO
	 * @throws Exception 
	 */
	public void TrainNextGeneration(Vector<double[]> trainingExamples, PrintStream outStream, boolean isLastStep) throws Exception
	{
		calculateFitness(trainingExamples, outStream, isLastStep);
		
		tournamentSelection(this.mutationProbability);	
		
	}

	private void tournamentSelection(double mutationProb) throws Exception 
	{
		nextGeneration = new ArrayList<Tree>();
		
		addElitistsFromCurrentGen();
		
		addNewRandomTreesToNextGen();
		
		while(nextGeneration.size() < this.populationSize)
		{
			Tree parent1 = selectParent();
			Tree parent2 = selectParent();
			
			Tree child = Tree.recombine(parent1, parent2);
			
			if (Utils.rand.nextDouble() < mutationProb)
			{
				child.mutate();
			}
			
			nextGeneration.add(child);			
		}
		
		currentGeneration = nextGeneration;
		
	}

	private void addElitistsFromCurrentGen() 
	{	
		if (elitism > 0)
		{
			SortedSet<Tree> bestTrees = new TreeSet<Tree>(new Comparator<Tree>() {

				@Override
				public int compare(Tree arg0, Tree arg1) 
				{
					if(Double.isInfinite(arg0.getFitness()) && Double.isInfinite(arg1.getFitness()))
					{
						return 0;
					}
					if (arg0.getFitness() > arg1.getFitness() || Double.isInfinite(arg0.getFitness()))
					{
						return -1; 
					}
					else if (arg0.getFitness() < arg1.getFitness() || Double.isInfinite(arg1.getFitness()))
					{
						return 1; 
					}
					return 0;
				}
			});

			for (Tree tree : currentGeneration)
			{
				if (!Double.isInfinite(tree.getFitness()) && tree.getFitness() < this.maxFitnessForElitism)
				{
					bestTrees.add(tree);
					if (bestTrees.size() > this.elitism)
					{
						bestTrees.remove(bestTrees.first());
					}
				}
			}
			
			for (Tree tree : bestTrees)
			{
				nextGeneration.add(tree);
			}
		}

	}

	private void addNewRandomTreesToNextGen() {
		//adding random new trees in each generation
		for (int i = 0; i < numRandomTreePerGen; i++)
		{
			nextGeneration.add(new Tree(maxDepth));
		}
	}

	private Tree selectParent() 
	{
		double bestFitness = Double.MAX_VALUE;
		Tree bestCandidate = null;
		
		//low fitness is better
		for(int i = 0; i < this.numTreesForTournamentSelection; i++)
		{
			Tree candidate = currentGeneration.get(Utils.rand.nextInt(currentGeneration.size()));
				
			if ((candidate.getFitness() < bestFitness) || (bestCandidate == null))
			{
				bestCandidate = candidate;
				bestFitness = candidate.getFitness();
			}
		}
		
		return bestCandidate;
	}

	private void calculateFitness(Vector<double[]> trainingExamples, PrintStream outStream, boolean isLastStep) {
		
		Tree bestTree = null;
		bestFitness = Double.MAX_VALUE;
		averageFitness = 0;
		
		double error;
		
		for (Tree tree : currentGeneration)
		{
			double fitness = 0;
			
			for (int i = 0; i < trainingExamples.size(); i++)
			{
				error = tree.evaluate(trainingExamples.get(i)[0]) - trainingExamples.get(i)[1];
				fitness += Math.abs(error);
			}
						
			fitness /= trainingExamples.size();
			int numNodes = tree.getNumNodes();
			double depthFactor = (numNodes > numNodesPenalty) ? ((fitness / 5.0)* numNodes + numNodes) : 0;
			fitness += depthFactor;
			
			if (Double.isNaN(fitness))
			{
				fitness = Double.MAX_VALUE;
			}
			tree.setFitness(fitness);
			
			averageFitness += fitness;
			
			if (fitness < bestFitness)
			{
				bestFitness = fitness;
				bestTree = tree;
			}
		}
		
		if (isLastStep)
		{
			//String msg = "\"" + bestFitness + "\",\"" + averageFitness + "\",\"" + bestTree + "\"";
			System.out.println(bestTree.toString());
			outStream.println(bestTree.toString());
		}

	}
	
	
	
}
